# -- coding: utf-8 --
"""
# MIAHX© 用户感知计算型传感器数据帧消息的 FastDDS 发布/订阅器

**模块：MIAHX© / 框架基础服务层（FFST） / 交互框架信息表示模型（IFIRM）**

Version: 0.4.10,
Last updated: October 20, 2024

## Contributor(s)
* YAO Naiming <[naiming@iscas.ac.cn](mailto:naiming@iscas.ac.cn)>

## Change log
### v0.4.10
- .20: [优化] Reader 构造方法传入 topic_name
- .19: [新增] 初始实现

---
Copyright © 2024-2025 MIAHX Contributors, ISCAS, and ASIC. All Rights Reserved.
Licensed under Apache-2.0.
"""

from enum import Enum

import miahx_fastdds as dds
import numpy as np

from .. import Participant
from ..reader import Reader, ReaderListener
from ..writer import Writer


class EDrivingGoalType(Enum):
    """
    驾驶人意图类型枚举
    """
    UNKNOWN = dds.DIT_UNKNOWN
    NON_STOP = dds.DIT_NON_STOP
    CHANGE_LANE_LEFT = dds.DIT_CHANGE_LANE_LEFT
    CHANGE_LANE_RIGHT = dds.DIT_CHANGE_LANE_RIGHT
    TURN_LEFT = dds.DIT_TURN_LEFT
    TURN_RIGHT = dds.DIT_TURN_RIGHT


class EUserKeypointType(Enum):
    """
    用户关键点类型枚举
    """
    IGNORE = dds.UKT_IGNORE
    FACE = dds.UKT_FACE
    BODY = dds.UKT_BODY
    LEFT_HAND = dds.UKT_LEFT_HAND
    RIGHT_HAND = dds.UKT_RIGHT_HAND


class UserStreamReader(Reader):
    """
    生理样本数据流读取器
    """

    def __init__(self, participant: Participant, topic_name: str, on_data_listeners=None):
        """
        创建生理样本数据流读取器

        :param participant: FastDDS 参与者
        :param topic_name: FastDDS 主题名称
        :param on_data_listeners: 数据监听回调列表
        """
        topic_data_type = dds.UserDataFramePubSubType()
        super().__init__(participant, topic_name, topic_data_type,
                         listener=ReaderListener(topic_name, topic_data_type, on_data_listeners))


class UserStreamWriter(Writer):
    """
    生理样本数据流写入器
    """

    def __init__(self, participant: Participant, topic_name: str):
        """
        创建用户感知样本数据流写入器

        :param participant: FastDDS 参与者
        :param topic_name: FastDDS 主题名称
        """
        super().__init__(participant, topic_name, dds.UserDataFramePubSubType())

    def write(self, chunk: list[tuple[int, float, EDrivingGoalType, np.ndarray]],
              iid: int = None, uts: float = None):
        """
        写入用户感知数据样本

        :param chunk: 样本数据描述，列表，每个元素为一个元组，表示：
        (样本的全局校准单调递增序号，样本的全局校准 POSIX 时间戳，驾驶人意图类型，用户关键点)
        :param iid: 全局校准的单调递增序号
        :param uts: 全局校准的 POSIX 时间戳
        """
        # header
        frame = dds.UserDataFrame()
        frame.header(self._create_frame_header(uts, iid))
        # body
        samples = []
        for item in chunk:
            iid, uts, ditype, kpts = item
            sample = dds.UserSample()
            sample.iid(iid)
            sample.uts(uts)
            if ditype is not None:
                sample.ditype(ditype.value)
            if kpts is not None:
                kpt_dims = kpts.shape[-1] - 1
                assert kpt_dims == 2 or kpt_dims == 3, 'kpt_dims != 2 or 3'
                sample.kpt_dims(kpt_dims)
                assert kpts.dtype == np.float32, 'kpts.dtype != np.float32'
                sample.keypoints(kpts.flatten().tolist())
            samples.append(sample)
        frame.samples(samples)
        # send
        self._writer.write(frame)
